import os
import sys
import random
import numpy as np
from scipy.io import loadmat
import matplotlib.pyplot as plt

#定义训练集和测试集的路径
TrainPath='statlearning-sjtu-2020/ECGTrainData/Train/'
TestPath='statlearning-sjtu-2020/ECGTestData/ECGTestData/'

#定义训练集和测试集的扩展长度
testDestLen=10120
trainDestLen=1256

#定义验证集合的比率
validRatio=0.2

cateMapping={'窦性心律_左室肥厚伴劳损':0,
'窦性心动过缓':1,
'窦性心动过速':2,
'窦性心律_不完全性右束支传导阻滞':3,
'窦性心律_电轴左偏':4,
'窦性心律_提示左室肥厚':5,
'窦性心律_完全性右束支传导阻滞':6,
'窦性心律_完全性左束支传导阻滞':7,
'窦性心律_左前分支阻滞':8,
'正常心电图':9}

def count_elements(seq)->dict:
    hist={}
    for x in seq:
        hist[x]=hist.get(x,0)+1
    return hist

def ascii_histogram(seq) -> None:
    counted = count_elements(seq)
    for k in sorted(counted):
        print('{0:5d} {1}'.format(k, '+' * counted[k]))

def trainLenAna():
    minLen=sys.maxsize;maxLen=0;lenList=[]
    folders=os.listdir(TrainPath)
    for folder in folders:
        folderPath=TrainPath+folder+'/'
        files=os.listdir(folderPath)
        for file in files:
            filePath=folderPath+file
            m=loadmat(filePath)
            beats=m['Beats'][0][0]
            beatData=beats[3][0]
            for i in range(len(beatData)):
                samplen=len(beatData[i])
                lenList.append(samplen)
                minLen=min(minLen,samplen)
                maxLen=max(maxLen,samplen)
    return minLen,maxLen,lenList

def testLenAna():
    minLen=sys.maxsize;maxLen=0;lenList=[]
    files=os.listdir(TestPath)
    for file in files:
        filePath=TestPath+file
        m=loadmat(filePath)
        samplen=len(m['data'])
        lenList.append(samplen)
        minLen=min(minLen,samplen)
        maxLen=max(maxLen,samplen)
    return minLen,maxLen,lenList

def drawDistribution(lenList,xLabel,yLabel,drawTitle,graphLabel):
    keyList=[]
    valList=[]
    lendict=count_elements(lenList)
    print(lendict)
    for key in sorted(lendict.keys()):
        keyList.append(key)
        valList.append(lendict[key])
    plt.bar(keyList,valList,label=graphLabel,width=4)
    plt.legend()
    plt.xlabel(xLabel)
    plt.ylabel(yLabel)
    plt.title(drawTitle)
    plt.show()

def getTest(index=6+1,destLen=trainDestLen):
    testData=[]
    fileNames=[]
    files=os.listdir(TestPath)
    for file in files:
        fileNames.append(file.split('.')[0])
        filePath=TestPath+file
        m=loadmat(filePath)
        data=m['data'].T[index][:destLen]
        testData.append(data)
    return fileNames,testData

def getOriginalTest(index=6+1):
    testData = []
    fileNames = []
    files = os.listdir(TestPath)
    for file in files:
        fileNames.append(file.split('.')[0])
        filePath = TestPath + file
        m = loadmat(filePath)
        data = m['data'].T[index]
        testData.append(data)
    return fileNames, testData


def getTrainValid(index=6,destLen=trainDestLen):
    trainData=[];trainLabels=[]
    validData=[];validLabels=[]
    folders = os.listdir(TrainPath)
    for folder in folders:
        cateList=[]
        label=cateMapping[folder]
        folderPath = TrainPath + folder + '/'
        files = os.listdir(folderPath)
        for file in files:
            filePath = folderPath + file
            m = loadmat(filePath)
            beats = m['Beats'][0][0]
            beatData = beats[3][0]
            for i in range(len(beatData)):
                sample=beatData[i].T
                cateList.append(np.resize(sample[index],destLen))
        catNum=len(cateList)
        validNum=int(catNum*validRatio)
        random.shuffle(cateList)
        validData.extend(cateList[:validNum])
        trainData.extend(cateList[validNum:])
        validLabels.extend([label for i in range(validNum)])
        trainLabels.extend([label for i in range(catNum-validNum)])
    return trainData,trainLabels,validData,validLabels




'''
minLen,maxLen,lenList=trainLenAna()
#minLen,maxLen,lenList=testLenAna()
print(minLen,maxLen)
print(lenList)
drawDistribution(lenList,'length','count','Test Length Distribution','Test Length Distribution')
'''
